import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import argparse
import os
import seaborn as sns

def parse_args():
    parser = argparse.ArgumentParser(description="Generate Scatter Quadrant Plot (Accuracy vs Precision).")
    parser.add_argument('--files','-f', nargs='+', required=True, help='List of files (CSV or Excel) to analyze')
    parser.add_argument('--exclude', nargs='+', default=['BART'], help='Methods to exclude (default: BART)')
    parser.add_argument('--output','-o', default='scatter_quadrant.png', help='Output filename for the plot')
    return parser.parse_args()



def read_data(file_path):
    _, ext = os.path.splitext(file_path)
    try:
        if ext.lower() in ['.xlsx', '.xls']:
            return pd.read_excel(file_path)
        else:
            return pd.read_csv(file_path)
    except Exception as e:
        print(f"Error reading {file_path}: {e}")
        return None

def get_method_stats(files, exclude_methods):
    """
    读取并聚合数据，计算每个方法的均值和标准差。
    """
    all_data = []
    
    for f in files:
        df = read_data(f)
        if df is not None and not df.empty:
            # 1. 过滤方法
            df = df[~df['method'].isin(exclude_methods)].copy()
            
            # 2. 重命名
            df['method'] = df['method'].replace('ours(MTM)', 'ours(deepsignal3)')
            
            # 3. 过滤Coverage (保留10的倍数)
            if 'coverage' in df.columns:
                df['coverage'] = pd.to_numeric(df['coverage'], errors='coerce')
                df = df.dropna(subset=['coverage'])
                df = df[np.isclose(df['coverage'] % 10, 0)]
            
            cols = ['method', 'pearson', 'RMSE']
            if all(c in df.columns for c in cols):
                all_data.append(df[cols])
            
    if not all_data:
        return None
        
    combined_df = pd.concat(all_data, ignore_index=True)
    
    # 计算均值和标准差
    stats = combined_df.groupby('method').agg(['mean', 'std']).reset_index()
    stats.columns = ['method', 'pearson_mean', 'pearson_std', 'RMSE_mean', 'RMSE_std']
    
    return stats

def plot_quadrant(stats_df, output_file):
    # 设置风格
    sns.set_style("whitegrid")
    sns.set_context("talk") 
    
    fig, ax = plt.subplots(figsize=(10, 8))
    
    # === 样式映射 ===
    style_map = {
        'Dorado': {'color': '#1f77b4', 'marker': 'o', 'label': 'Dorado (Gold Standard)'},
        'ours(deepsignal3)': {'color': '#d62728', 'marker': 'D', 'label': 'Ours (deepsignal3)'}, # 红色菱形
        'Rockfish': {'color': '#2ca02c', 'marker': '^', 'label': 'Rockfish'},
        'DeepMod2': {'color': '#ff7f0e', 'marker': 's', 'label': 'DeepMod2'}
    }
    
    # === 绘制散点和误差线 ===
    for _, row in stats_df.iterrows():
        method = row['method']
        x = row['pearson_mean']
        y = row['RMSE_mean']
        xerr = row['pearson_std']
        yerr = row['RMSE_std']
        
        style = style_map.get(method, {'color': 'gray', 'marker': 'o', 'label': method})
        
        # 绘制带误差线的点 (移除 ax.annotate)
        ax.errorbar(x, y, xerr=xerr, yerr=yerr, 
                    fmt=style['marker'], 
                    color=style['color'], 
                    ecolor='lightgray', 
                    elinewidth=2, 
                    capsize=5, 
                    markersize=15, 
                    label=style['label'],
                    markeredgecolor='white',
                    markeredgewidth=1.5,
                    alpha=0.9)

    # === 坐标轴与标签 ===
    ax.set_xlabel('Pearson Correlation (Higher is Better)', fontweight='bold')
    ax.set_ylabel('RMSE (Lower is Better)', fontweight='bold')
    ax.set_title('Accuracy vs. Precision Trade-off', fontweight='bold', pad=20)
    
    # === 优化显示范围 ===
    x_min, x_max = stats_df['pearson_mean'].min(), stats_df['pearson_mean'].max()
    y_min, y_max = stats_df['RMSE_mean'].min(), stats_df['RMSE_mean'].max()
    
    margin_x = (x_max - x_min) * 0.2
    margin_y = (y_max - y_min) * 0.2
    
    ax.set_xlim(x_min - margin_x, x_max + margin_x)
    ax.set_ylim(y_min - margin_y, y_max + margin_y)

    # === 绘制背景“最佳区域”指示箭头 ===
    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    
    # 箭头指向右下角
    arrow_start_x = xlim[0] + (xlim[1]-xlim[0])*0.1
    arrow_start_y = ylim[1] - (ylim[1]-ylim[0])*0.1
    arrow_end_x = xlim[1] - (xlim[1]-xlim[0])*0.1
    arrow_end_y = ylim[0] + (ylim[1]-ylim[0])*0.1
    
    ax.annotate("", xy=(arrow_end_x, arrow_end_y), xytext=(arrow_start_x, arrow_start_y),
                arrowprops=dict(arrowstyle="->", color="0.9", lw=5))
    
    ax.text(arrow_end_x, arrow_end_y + (ylim[1]-ylim[0])*0.05, "", 
            ha='right', va='bottom', fontsize=12, color='gray', fontweight='bold')

    # 图例 (Legend) - 这是现在唯一的识别方式
    plt.legend(loc='upper right', frameon=True, framealpha=0.9, fontsize=11)
    
    sns.despine()
    plt.tight_layout()
    
    print(f"Saving cleaned scatter plot to {output_file}")
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    plt.show()

def main():
    args = parse_args()
    
    print("Reading and aggregating data...")
    stats = get_method_stats(args.files, args.exclude)
    
    if stats is not None:
        plot_quadrant(stats, args.output)
    else:
        print("No valid data found.")

if __name__ == "__main__":
    main()